#pragma once
#include "task.h"
#include <concepts>
#include <cstddef>
#include <cstring>
#include <iostream>
#include <optional>
#include <span>
#include <string_view>
#include <utility>
#include <vector>

namespace co_async {

template <class Reader> class IStreamBase {
public:
  explicit IStreamBase(std::size_t cap = 8192) : m_buffer(cap) {}

  IStreamBase(IStreamBase<Reader> &&that) = default;
  IStreamBase &operator=(IStreamBase &&) = default;

  Task<std::optional<char>> getchar() {
    if (m_index == m_size) {
      if (!co_await read_all()) {
        co_return std::nullopt;
      }
    }
    co_return m_buffer[m_index++];
  }

  Task<std::string> getline(int end = '\n') {
    std::string line;
    while (true) {
      for (m_index; m_index < m_size; m_index++) {
        auto c = m_buffer[m_index];
        if (c == end) {
          m_index++;
          co_return line;
        }
        line.push_back(c);
      }
      if (!co_await read_all()) {
        co_return line;
      };
    }
  }

  Task<std::string> getline(std::string_view eol) {
    std::string line = "";
    std::size_t target_idx = 0;
    std::size_t eol_size = eol.size();
    while (true) {
      for (; m_index < m_size; m_index++) {
        auto c = m_buffer[m_index];
        if (c == eol[target_idx]) {
          target_idx++;
          if (target_idx == eol_size) {
            m_index++;
            co_return line;
          }
          continue;
        }
        line.push_back(c);
        target_idx = 0;
      }
      if (!co_await read_all()) {
        co_return line;
      };
    }
  }

  Task<std::string> getn(std::size_t n) {
    std::string line(n, 0);
    std::size_t already_copy = 0;
    while (true) {
      std::size_t avali_size = m_size - m_index;
      std::size_t not_copy = n - already_copy;
      if (avali_size >= not_copy) {
        line.replace(already_copy, not_copy, m_buffer.data() + m_index,
                     not_copy);
        m_index += not_copy;
        co_return line;
      }
      line.replace(already_copy, avali_size, m_buffer.data() + m_index,
                   avali_size);
      already_copy += avali_size;
      if (!co_await read_all()) {
        line.resize(already_copy);
        co_return line;
      }
    }
    co_return line;
  }

private:
  Task<bool> read_all() {
    m_size = co_await static_cast<Reader *>(this)->read(m_buffer);
    m_index = 0;
    co_return m_size;
  }
  std::vector<char> m_buffer;
  int fd;
  std::size_t m_index = 0;
  std::size_t m_cap = 0;
  std::size_t m_size = 0;
};

template <class Writer> class OStreamBase {
public:
  explicit OStreamBase(std::size_t cap = 8192)
      : m_index(0), m_buffer(cap), m_cap(cap) {}

  OStreamBase(OStreamBase<Writer> &&that) = default;
  OStreamBase &operator=(OStreamBase &&) = default;

  Task<> putchar(char c) {
    if (m_index == m_cap) {
      co_await flush();
      co_return;
    }
    m_buffer[m_index++] = c;
  }

  Task<> puts(std::string_view s) {
    std::size_t already_put = 0;
    for (auto c : s) {
      if (m_index == m_cap) {
        co_await flush();
        m_index = 0;
      }
      m_buffer[m_index++] = c;
    }
  }

  Task<> putline(std::string_view s) {
    co_await puts(s);
    co_await putchar('\n');
    co_await flush();
  }

  struct ShutdownException {};

  Task<> flush() {
    if (!m_index) [[unlikely]] {
      co_return;
    }
    Writer *that = static_cast<Writer *>(this);
    auto buf = std::span<char>(m_buffer.data(), m_index);
    std::size_t already_write = 0;
    std::size_t len = 0;
    while (already_write < m_index) {
      len = co_await that->write(buf);
      buf = buf.subspan(len);
      if (len <= 0) [[unlikely]] {
        m_index = 0;
        throw ShutdownException{};
      }
      already_write += len;
    }
    m_index = 0;
  }

private:
  std::vector<char> m_buffer;
  std::size_t m_index;
  std::size_t m_cap;
};

template <class StreamBuf>
class IOStreamBase : public IStreamBase<StreamBuf>,
                     public OStreamBase<StreamBuf> {
public:
  explicit IOStreamBase(std::size_t cap = 65536)
      : IStreamBase<StreamBuf>(cap), OStreamBase<StreamBuf>(cap) {}
  
  IOStreamBase(IOStreamBase &&that) = default;
  IOStreamBase &operator=(IOStreamBase &&) = default;
  
};

template <class StreamBuf>
class [[nodiscard]] IOStream : public IOStreamBase<IOStream<StreamBuf>>,
                               public StreamBuf {
public:
  template <class... Args>
    requires std::constructible_from<StreamBuf, Args...>
  explicit IOStream(Args &&...args)
      : IOStreamBase<IOStream<StreamBuf>>(),
        StreamBuf(std::forward<Args>(args)...) {}
  IOStream() = default;
  IOStream(IOStream &&that) = default;
  IOStream &operator=(IOStream &&) = default;
};

template <class StreamBuf>
class [[nodiscard]] OStream : public OStreamBase<OStream<StreamBuf>>,
                              public StreamBuf {
public:
  template <class... Args>
    requires std::constructible_from<StreamBuf, Args...>
  explicit OStream(Args &&...args)
      : OStreamBase<StreamBuf>(), StreamBuf(std::forward<Args>(args)...) {}
  OStream() = default;
  OStream(OStream &&that) = default;
  OStream &operator=(OStream &&) = default;
};

template <class StreamBuf>
class [[nodiscard]] IStream : public IStreamBase<IStream<StreamBuf>>,
                              public StreamBuf {
public:
  template <class... Args>
    requires std::constructible_from<StreamBuf, Args...>
  explicit IStream(Args &&...args)
      : IStreamBase<StreamBuf>(), StreamBuf(std::forward<Args>(args)...) {}
  IStream() = default;

  IStream(IStream &&that) = default;
  IStream &operator=(IStream &&) = default;

};

} // namespace co_async